# -*- coding: utf-8 -*-

# 作用：采用Python进行曲线拟合
# ref: https://www.geekering.com/categories/machine-learning/bruno-silva/
# machine-learning-python-non-linear-regression/

# -----------------------------
# 导入函数库
# -----------------------------

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from sklearn.metrics import r2_score

# -----------------------------
# 加载和显示数据
# -----------------------------

data_dir = './data_files/'
data_file = data_dir + "gdp_since_1960.csv"
dataset = pd.read_csv(data_file)
num_dataset = len(dataset)
print("数据长度： " + str(num_dataset))
print("前10个数据为：")
print(dataset.head(10))

# 数据画图，判断大致趋势

plt.figure(figsize=(8,5))
x_data = dataset["Year"].values
y_data = dataset["Value"].values
plt.plot(x_data, y_data, 'ro')
plt.ylabel('GDP')
plt.xlabel('Year')
plt.show()

# 数据归一化

x_data_norm = x_data / max(x_data)
y_data_norm = y_data / max(y_data)

# 随机挑出20%数据作为测试数据

np.random.seed(4711) # 确保每次随机出来的数是一样的
rand_mask = np.random.rand(num_dataset) < 0.85
train_x = x_data_norm[rand_mask]
test_x = x_data_norm[~rand_mask]
train_y = y_data_norm[rand_mask]
test_y = y_data_norm[~rand_mask]

plt.figure(figsize=(8,5))
plt.plot(train_x, train_y, 'ro', label='Train Data')
plt.plot(test_x, test_y, 'bs', label='Test Data')
plt.ylabel('GDP (Scaled)')
plt.xlabel('Year (Scaled)')
plt.legend(loc='best')
plt.show()

# -----------------------------
# 函数模型
# -----------------------------

# 定义函数模型
def sigmoid(x, beta1, beta2):
    y = 1.0 / (1.0 + np.exp(-beta1 * (x - beta2)))
    return y


# 函数模型画图

x_model = np.arange(-5.0, 5.0, 0.1)
y_model = sigmoid(x_model, 1.0, 0.0)
plt.plot(x_model, y_model)
plt.ylabel('Dependent Variable')
plt.xlabel('Independent Variable')
plt.show()

# -----------------------------
# 拟合参数
# -----------------------------

popt, pcov= curve_fit(sigmoid, train_x, train_y,
                      method='lm')
beta1 = popt[0]
beta2 = popt[1]

# 输出最终参数

print("beta1 = %f, beta2 = %f" % (beta1, beta2))

# 画出拟合结果

x_fit = np.linspace(1960, 2015, 55)
x_fit = x_fit / max(x_fit)
y_fit = sigmoid(x_fit, beta1, beta2)
plt.plot(train_x, train_y, 'ro', label='Train Data')
plt.plot(test_x, test_y, 'bs', label='Test Data')
plt.plot(x_fit, y_fit, linewidth=3.0, label='Fit Curve')
plt.legend(loc='best')
plt.ylabel('GDP')
plt.xlabel('Year')
plt.show()

# 损失

y_loss = sigmoid(train_x, beta1, beta2)
loss_abs = np.mean((y_loss - train_y) ** 2)
print("拟合损失：%.4f" % (loss_abs, ))

# 通过测试数据查看模型的准确度

y_hat = sigmoid(test_x, beta1, beta2)
residual_sum_sq = np.mean((y_hat - test_y) ** 2)
test_score = r2_score(y_hat , test_y)
print("残差平方和：%.4f" % (residual_sum_sq))
print("R2评分：%.4f" % (test_score))

# -----------------------------
# 使用训练好的模型
# -----------------------------

x_new = np.linspace(2015, 2020, 5)
x_new = x_new / max(x_data)
y_new = sigmoid(x_new, beta1, beta2)

plt.plot(train_x, train_y, 'ro', label='Train Data')
plt.plot(x_new, y_new, 'ks', label='Prediction')
plt.plot(x_fit, y_fit, linewidth=3.0, label='Fit Curve')
plt.legend(loc='best')
plt.ylabel('GDP')
plt.xlabel('Year')
plt.show()

# -----------------------------
# 保存和加载训练好的模型
# -----------------------------

with open(data_dir + 'fitting_result.txt', 'w') as frst:
    frst.write("beta1 = %.4f\n" % (beta1, ))
    frst.write("beta1 = %.4f\n" % (beta2, ))
